#!/usr/bin/env python3
# Author: Armit
# Create Time: 2022/11/18 

import os
from glob import glob
from pathlib import Path
from collections import Counter
from argparse import ArgumentParser

import numpy as np
import pandas as pd
import matplotlib.pylab as plt
from moviepy.editor import ImageSequenceClip
from sklearnex import patch_sklearn ; patch_sklearn()
from sklearn.neighbors import KNeighborsClassifier
from sklearn.manifold import TSNE

from data import get_data, FEATURE_NUM, cat_dict, LOG_PATH
from utils import RAND_SEED, CPU_COUNT, get_cmap
from pca import _pca, _display

Path(LOG_PATH).mkdir(exist_ok=True)

TARGET = 'Prop3'
UNDEFINED_ID = cat_dict.get_cat_id(TARGET, 'Undefined')
print(f'UNDEFINED_ID: {UNDEFINED_ID}')


def _tsne_fast(X:pd.DataFrame, dim=2) -> pd.DataFrame:
  tsne = TSNE(n_components=dim, n_iter=300, learning_rate='auto', init='random', random_state=RAND_SEED, n_jobs=CPU_COUNT, square_distances=True)
  X_hat = tsne.fit_transform(X)
  print(f'  kl_divergence: {tsne.kl_divergence_}')
  return X_hat


def make_stat(y_pred):
  cntr = Counter(y_pred)
  for id in sorted(cntr.keys()):
    print(f'  assign {cntr[id]} ({cntr[id] / len(y_pred):.3%}) items to class {cat_dict.get_cat_name(TARGET, id)}')


def make_plot(X:np.ndarray, y:np.array, y_hat:np.array, name='', maximize=True):
  # y: original 4-class labels, y_hat: (partially) refilled 3-class labels
  assert len(X) == len(y) == len(y_hat)

  # plot
  projector = _tsne_fast if args.tsne else _pca
  X_hat = projector(X)
  x_min, x_max = X_hat[:, 0].min(), X_hat[:, 0].max()
  y_min, y_max = X_hat[:, 1].min(), X_hat[:, 1].max()

  plt.clf()
  plt.subplot(121)
  plt.scatter(X_hat[:, 0], X_hat[:, 1], s=1, c=y, cmap=get_cmap(y))
  plt.xlim(x_min, x_max) ; plt.ylim(y_min, y_max)
  plt.title(f'{TARGET} original')
  plt.subplot(122)
  plt.scatter(X_hat[:, 0], X_hat[:, 1], s=1, c=y_hat, cmap=get_cmap(y_hat))
  plt.xlim(x_min, x_max) ; plt.ylim(y_min, y_max)
  plt.title(f'{TARGET} after reassignment')
  plt.tight_layout()
  plt.savefig(Path(LOG_PATH) / f'{args.prefix}{name}.png')
  if args.show:
    if maximize:
      backend = plt.get_backend()
      mng = plt.get_current_fig_manager()
      if   backend == 'QtAgg': mng.window.showMaximized()
      elif backend == 'wxAgg': mng.frame.Maximize(True)
      elif backend == 'TkAgg': mng.window.state('zoomed')
      else: print(f'unknow backend: {backend}')
    plt.show()


def knn_iterative(X_train, y_train, X_test, k, p, ratio=0.1):
  N = len(X_test)
  r = int(N * ratio)      # 每次只指派平均近邻距离最小的前 r 个未知点

  i = 0
  while len(X_test):
    i += 1
    print(f'>> [step {i}] remaining: {len(X_test)}, total: {N}')

    # knn
    knn = KNeighborsClassifier(n_neighbors=k, p=p, weights='distance', n_jobs=os.cpu_count())
    knn.fit(X_train, y_train)

    # filter top-r avg-nearest samples
    _, dists = knn.kneighbors(X_test)       # [N, K]
    avg_dists = dists.mean(axis=-1)         # [N]
    sorted_index = avg_dists.argsort()
    select_index = sorted_index[:r]         # [R]

    # make assignment
    X_test_r = X_test.iloc[select_index]    # [R, D]
    y_test_r = knn.predict(X_test_r)        # [R]
    index = X_test_r.index
    X_test = X_test.drop(index)

    # merge (semi-supervise)
    X_train = pd.concat([X_train, X_test_r], axis=0, copy=False)
    y_orig  = pd.concat([y_train, pd.Series(data=np.ones_like(y_test_r) * UNDEFINED_ID, index=index)], axis=0, copy=False)
    y_train = pd.concat([y_train, pd.Series(data=y_test_r, index=index)], axis=0, copy=False)

    # plot midway results
    make_stat(y_test_r)
    make_plot(X_train, y_orig, y_train, name=f'knn_color-{i:04d}')

  # save gif
  files = glob(str(Path(LOG_PATH) / f'{args.prefix}knn_color-*.png'))
  files += [files[-1]] * args.fps    # stall ending for 1s
  clip = ImageSequenceClip(files, fps=args.fps)
  clip.write_gif(Path(LOG_PATH) / f'{args.prefix}knn_color-iterative.gif')

  # show 3d view
  if args.show_3d:
    X_hat = _pca(X_train, 'tsne', dim=3)
    _display(X_hat, y_train, dim=3, title=TARGET)


def knn_color(args):
  X, y = get_data(args.limit, FEATURE_NUM, target=TARGET, sample_method='same')
  mask = y == UNDEFINED_ID
  X_train, y_train = X[~mask], y[~mask]
  X_test,  y_test  = X[ mask], y[ mask]
  print(f'dataset: {len(X_train)} for train. {len(X_test)} samples for test')
  
  print(f'[kNN] k={args.k} p={args.p}')
  if args.baseline:
    knn = KNeighborsClassifier(n_neighbors=args.k, p=args.p, weights='distance', n_jobs=4)
    knn.fit(X_train, y_train)

    y_pred = knn.predict(X_test)
    make_stat(y_pred)

    # back-fill assignment
    y_orig = y.tolist()
    j = 0
    y_hat = np.empty_like(y)
    for i in range(len(y_hat)):
      if y_orig[i] == UNDEFINED_ID:
        y_hat[i] = y_pred[j]
        j += 1
      else:
        y_hat[i] = y_orig[i]
    # assure all filled
    assert j == len(y_pred)
    assert len(set(y))     == cat_dict.get_cat_ord(TARGET)
    assert len(set(y_hat)) == cat_dict.get_cat_ord(TARGET) - 1

    # plot
    make_plot(X, y, y_hat, name='knn_color')

    # show 3d view
    if args.show_3d:
      X_hat = _pca(X, 'tsne', dim=3)
      _display(X_hat, y_hat, dim=3, title=TARGET)
  else:
    knn_iterative(X_train, y_train, X_test, args.k, args.p, args.r)


if __name__ == '__main__':
  parser = ArgumentParser(description='Semi-Superivised kNN classification')
  parser.add_argument('-k', default=5, type=int, help='kNN n_neighbors')
  parser.add_argument('-p', default=2, type=int, help='minkowski distance power')
  parser.add_argument('-r', default=0.01, type=float, help='ratio in iterative method')
  parser.add_argument('-N', '--limit', default=40000, type=int, help='limit dataset size')
  parser.add_argument('--fps', default=15, type=float, help='FPS of gif')
  parser.add_argument('--baseline', action='store_true', help='non-iterative knn_color (baseline)')
  parser.add_argument('--tsne', action='store_true', help='project by T-SNE instead of PCA (NOTE: very slow, data points not stable in figure)')
  parser.add_argument('--show', action='store_true', help='show 2d-projected figure')
  parser.add_argument('--show_3d', action='store_true', help='show final results in 3d view')
  args = parser.parse_args()

  args.prefix = 'tsne_' if args.tsne else ''

  knn_color(args)
